from typing import Any
from sentence_transformers import SentenceTransformer
from chromadb import Documents, EmbeddingFunction, Embeddings

class Embeddings(EmbeddingFunction):
    def __init__(self, model_path):
        self.model = SentenceTransformer(model_path)

    def __call__(self, texts: Documents) -> Embeddings:
        # embed the documents somehow
        embeddings = self.embed_documents(texts)
        return embeddings
    def embed_documents(self, text_list):
        embeddings = self.model.encode(text_list)
        encod_list = embeddings.tolist()
        return encod_list

    def embed_query(self, text):
        embeddings = self.model.encode([text])
        encod_list = embeddings.tolist()
        return encod_list[0]
    

if __name__ == "__main__":
    import os
    dir_name = os.path.dirname(__file__)
    model_path = os.path.join(dir_name, 'model/m3e')
    print(os.path.abspath(model_path))
    embed = Embeddings(model_path)
    print(embed.embed_query('你好呀'))
